from matplotlib import pyplot as plt
from sklearn.metrics import ConfusionMatrixDisplay, classification_report, confusion_matrix

import data
from train import load_model_weight
import numpy as np

test_data, test_label = data.test_data()
test_data = test_data.reshape(-1, 32, 32, 3) / 255.0
categories = [i.decode() for i in data.categories()]

# 加载数据
model = load_model_weight('./model/weights_with_data_enhancement.h5')
pred_label = model.predict(test_data)
pred_label = np.argmax(pred_label, axis=1)
print(pred_label)

cm = confusion_matrix(test_label, pred_label)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=categories)
fig, ax = plt.subplots(figsize=(10, 10))
disp = disp.plot(xticks_rotation='vertical', ax=ax, cmap='summer')

print(classification_report(test_label, pred_label))

plt.show()
